import os
import sys
import time
import warnings

import cv2
import numpy as np
import openpyxl
import pandas as pd
import torch
from PyQt5.QtCore import QThreadPool, QRunnable, pyqtSignal, QObject, QEventLoop
from PyQt5.QtCore import QTimer
from PyQt5.QtGui import QImage
from PyQt5.QtGui import QPixmap, QTransform
from PyQt5.QtWidgets import QApplication, QMainWindow
from PyQt5.QtWidgets import QFileDialog
from PyQt5.QtWidgets import QGraphicsPixmapItem
from PyQt5.QtWidgets import QGraphicsScene
from PyQt5.QtWidgets import QMessageBox
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.figure import Figure
from scipy.fft import fft

from UI import *
from detect import yolov8


class WorkerSignals(QObject):
    finished = pyqtSignal(object, object, object)


class YoloWorker(QRunnable):
    def __init__(self, yolov8_func, path, model, device, IoU, conf, save_path, image_name, csv_error):
        super(YoloWorker, self).__init__()
        self.yolov8_func = yolov8_func
        self.path = path
        self.model = model
        self.device = device
        self.IoU = IoU
        self.conf = conf
        self.save_path = save_path
        self.image_name = image_name
        self.csv_error = csv_error
        self.signals = WorkerSignals()

    def run(self):
        img, res = self.yolov8_func(self.path, self.model, self.device, self.IoU, self.conf,
                                    self.save_path, self.image_name, self.csv_error)
        self.signals.finished.emit(img, res, self.path)


class MyWindow(QMainWindow, Ui_GUI):
    def __init__(self, parent=None):
        super(MyWindow, self).__init__(parent)

        self.df = None
        self.error = False
        self.pic_path = None
        self.file_content = None
        self.threadpool = QThreadPool()
        self.results = {}
        self.img2show_result = None
        self.img2show_completed = False

        self.setupUi(self)

        self.load_data.clicked.connect(self.read_csv)  # 读取钢轨csv文件信息
        self.choose_speed.clicked.connect(self.read_speed)  # 读取钢轨csv文件信息
        self.load_pic.clicked.connect(self.read_photo)  # 读取钢轨图片信息

        self.go.clicked.connect(self.main)  #执行

        self.next_one.clicked.connect(self.next)  # 下一个
        self.last_one.clicked.connect(self.last)  # 上一个

        self.IoU_slider.valueChanged.connect(self.IoUslider_valuechange)  # 修改IoU的滑动条
        self.IoU_num.valueChanged.connect(self.IoUnum_valuechange)  # 修改IoU的数字条

        self.conf_slider.valueChanged.connect(self.confslider_valuechange)  # 修改conf的滑动条
        self.conf_num.valueChanged.connect(self.confnum_valuechange)  # 修改conf的数字条

        self.imgjpg_slider.valueChanged.connect(self.imgjpg_slider_valuechange)  # 修改图片范围的滑动条
        self.imgjpg_num.valueChanged.connect(self.imgjpg_num_valuechange)  # 修改图片范围的数字条

        self.timer = QTimer()
        self.timer.setSingleShot(True)  # 设置定时器为单次触发
        self.timer.timeout.connect(self.next)
        self.auto_next.toggled.connect(self.auto)
        self.timer = QTimer()
        self.timer.setInterval(2000)  # 设置定时器间隔为2秒
        self.timer.timeout.connect(self.next)

        self.channel1.stateChanged.connect(self.channel_show)
        self.channel2.stateChanged.connect(self.channel_show)
        self.channel3.stateChanged.connect(self.channel_show)
        self.channel4.stateChanged.connect(self.channel_show)

        self.pinpu.clicked.connect(self.show_waveform)

    def read_csv(self):
        file_dialog = QFileDialog()
        file_dialog.setNameFilter("CSV Files (*.csv)")
        file_dialog.exec_()
        selected_files = file_dialog.selectedFiles()
        if selected_files:
            self.load_text.setText(selected_files[0])
            # 获取文件的父路径
            parent_path = os.path.dirname(selected_files[0])
            filename = os.path.basename(selected_files[0])[0:-4]
            self.file_content = {
                "parent_path": parent_path,
                "filename": int(filename)
            }
            self.csv_num_2.setText(str(self.file_content["filename"]))  # 显示csv序号

    def read_speed(self):
        file_dialog = QFileDialog()
        file_dialog.setNameFilter("Excel Files (*.xlsx)")
        file_dialog.exec_()
        selected_files = file_dialog.selectedFiles()
        if selected_files:
            # 使用整个路径作为 filename
            filename = selected_files[0]
            self.speed_text.setText(filename)  # 显示文件路径
            self.df = pd.read_excel(filename)

    def read_photo(self):
        pic_path = QFileDialog.getExistingDirectory(self, '选择图片所在文件夹', "/", QFileDialog.ShowDirsOnly)
        if pic_path:
            self.pic_path = pic_path
            self.load_pho.setText(str(self.pic_path))  # 显示csv序号

    def get_data(self):  #获取振动数据
        file_path = self.check_path()
        if os.path.exists(file_path):
            self.load_text.setText(file_path)
            df = pd.read_csv(file_path)
            return df

    def running(self):
        try:
            start_time = time.perf_counter()
            self.error = False
            csv_path = self.check_path()  # 读取路径csv
            ku = self.channel_show()  # 显示振动频率
            self.save_ku(ku)  # 保存峭度
            self.img_kind()  # 处理图片
            end_time = time.perf_counter()
            running_time = end_time - start_time
            print("运行时间为：", running_time, "秒")
            if self.error:
                self.csv_alert()
            return csv_path
        except Exception as e:
            print("An error occurred:", e)
            csv_path = None
            return csv_path

    def save_ku(self, ku):
        for k in range(len(ku)):
            if float(ku[k]) > float(10):
                self.error = True
        ku_path = r"results/峭度.xlsx"
        # 尝试加载现有文件，如果不存在则创建一个新的
        try:
            workbook = openpyxl.load_workbook(ku_path)
        except FileNotFoundError:
            workbook = openpyxl.Workbook()

        # 选择工作表
        sheet = workbook.active

        # 检查 ku 的第一个元素是否在 Excel 的第一列中
        first_column_values = [sheet.cell(row=i, column=1).value for i in range(1, sheet.max_row + 1)]
        if self.file_content["filename"] not in first_column_values:
            # 如果不存在，则插入
            ku.insert(0, self.file_content["filename"])

            # 找到第一个空行
            row = sheet.max_row + 1
            # 将列表中的值追加到 Excel 中
            for col, value in enumerate(ku, start=1):
                sheet.cell(row=row, column=col, value=value)
            # 保存文件
            workbook.save(ku_path)

    def img_kind(self):
        pic_path20, pic_path21, image_name = self.check_photo()
        self.left_show.setScene(None)
        self.right_show.setScene(None)
        # 定义要查找的变量
        target_variable = image_name
        # 找到变量所在行的索引
        row_index = self.df[self.df.iloc[:, 0] == target_variable].index
        # 检查是否找到变量
        if not row_index.empty:
            # 获取该行第三列的信息
            third_column_value = self.df.iloc[row_index[0], 2]
            #print(third_column_value)
        else:
            third_column_value = 0
            print("变量未找到")
        self.speed.setText("当前速度：%s km/h" % third_column_value)
        self.img2show(pic_path20, pic_path21, image_name)

    def next(self):
        try:
            self.res_show.clear()
            self.file_content["filename"] = self.file_content["filename"] + 1
            csv_path = self.running()
            if not os.path.exists(csv_path):
                self.file_content["filename"] = self.file_content["filename"] - 1
        except:
            pass

    def last(self):
        try:
            self.res_show.clear()
            self.file_content["filename"] = self.file_content["filename"] - 1
            csv_path = self.running()
            if not os.path.exists(csv_path):
                self.file_content["filename"] = self.file_content["filename"] + 1
        except:
            pass

    def img2show(self, path1, path2, image_name):
        model = 'YOLOv8.pt'
        conf = self.conf_num.value()
        IoU = self.IoU_num.value()
        self.img2show_completed = False
        self.img2show_result = None

        if image_name > 0:

            image_name = str(image_name)

            save_path20 = r'results/{}/{}.jpg'.format(image_name, image_name + "_20")
            save_path21 = r'results/{}/{}.jpg'.format(image_name, image_name + "_21")
            csv_error = self.error

            if self.CUDA.isChecked():
                device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
            else:
                device = torch.device('cpu')
            print(device)
            if self.worker.isChecked():
                worker1 = YoloWorker(yolov8, path1, model, device, IoU, conf, save_path20, image_name, csv_error)
                worker2 = YoloWorker(yolov8, path2, model, device, IoU, conf, save_path21, image_name, csv_error)

                worker1.signals.finished.connect(self.handle_result)
                worker2.signals.finished.connect(self.handle_result)

                self.threadpool.start(worker1)
                self.threadpool.start(worker2)

                # 等待结果
                while not self.img2show_completed:
                    QApplication.processEvents()  # 保持GUI响应

            else:
                img1, res1 = yolov8(path1, model, device, IoU, conf, save_path20, image_name, csv_error)
                img2, res2 = yolov8(path2, model, device, IoU, conf, save_path21, image_name, csv_error)
                frame1 = self.RGB(img1)
                frame2 = self.RGB(img2)

                scene1 = self.img_transform(frame1)
                scene2 = self.img_transform(frame2)
                self.left_show.setScene(scene1)
                self.right_show.setScene(scene2)
                self.img2show_result = res1 == 1 or res2 == 1
        return self.img2show_result

    def handle_result(self, img, res, path):
        self.results[path] = (img, res)
        if len(self.results) == 2:
            self.process_results()

    def process_results(self):
        img1, res1 = self.results[list(self.results.keys())[0]]
        img2, res2 = self.results[list(self.results.keys())[1]]

        frame1 = self.RGB(img1)
        frame2 = self.RGB(img2)

        scene1 = self.img_transform(frame1)
        scene2 = self.img_transform(frame2)
        self.left_show.setScene(scene1)
        self.right_show.setScene(scene2)

        self.img2show_result = res1 == 1 or res2 == 1
        self.img2show_completed = True

        # 清理结果字典，为下一次操作做准备
        self.results.clear()

    def RGB(self, img):
        # 将BGR转换为RGB
        rgb_image = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

        # 创建QImage
        height, width, channel = rgb_image.shape
        bytes_per_line = 3 * width
        q_image = QImage(rgb_image.data, width, height, bytes_per_line, QImage.Format_RGB888)

        # 现在你可以使用这个q_image在PyQt界面中显示图像：
        # 例如，如果你有一个QLabel叫做image_label，你可以这样设置：
        # self.image_label.setPixmap(QPixmap.fromImage(q_image))：

        return q_image

    def img_transform(self, frame):
        # 生成qt能展示的图
        pix = QPixmap.fromImage(frame)
        transform = QTransform().rotate(90)
        pix = pix.transformed(transform)
        # 缩小图像
        height = 2.5
        weight = 0.5
        with warnings.catch_warnings():
            warnings.filterwarnings("ignore", category=DeprecationWarning)
            pix = pix.scaled(pix.width() * weight,
                             pix.height() * height,
                             aspectRatioMode=QtCore.Qt.AspectRatioMode.KeepAspectRatio)
        item = QGraphicsPixmapItem(pix)
        scene = QGraphicsScene()
        scene.addItem(item)
        return scene

    def channel_show(self):
        try:
            data = self.get_data()

            channel_1 = data.iloc[:, 2].values.astype(float)
            channel_2 = data.iloc[:, 5].values.astype(float)
            channel_3 = data.iloc[:, 8].values.astype(float)
            channel_4 = data.iloc[:, 11].values.astype(float)

            channel_1 = channel_1 * 1000 / (2.463 * 9.8)
            channel_2 = channel_2 * 1000 / (2.46 * 9.8)
            channel_3 = channel_3 * 1000 / (2.466 * 9.8)
            channel_4 = channel_4 * 1000 / (2.427 * 9.8)

            X_kf = []
            X_kf1 = self.features(channel_1)  # 计算峭度
            X_kf2 = self.features(channel_2)  # 计算峭度
            X_kf3 = self.features(channel_3)  # 计算峭度
            X_kf4 = self.features(channel_4)  # 计算峭度
            X_kf.append(X_kf1)
            X_kf.append(X_kf2)
            X_kf.append(X_kf3)
            X_kf.append(X_kf4)
            res = f"1L-Z峭度指标: {X_kf1}\n"
            res += f"1R-Z峭度指标: {X_kf2}\n"
            res += f"4L-Z峭度指标: {X_kf3}\n"
            res += f"4R-Z峭度指标: {X_kf4}\n"

            self.res_show.setText(res)

            channel_data = {
                1: channel_1,
                2: channel_2,
                3: channel_3,
                4: channel_4
            }

            self.channel.setScene(None)

            # 创建新的图形场景
            scene = QtWidgets.QGraphicsScene()
            self.channel.setScene(scene)

            # 创建新的图形和画布
            figure = Figure(figsize=(12, 2))
            canvas = FigureCanvas(figure)

            # 获取选中的通道
            selected_channels = []
            if self.channel1.isChecked():
                selected_channels.append(1)
            if self.channel2.isChecked():
                selected_channels.append(2)
            if self.channel3.isChecked():
                selected_channels.append(3)
            if self.channel4.isChecked():
                selected_channels.append(4)

            # 绘制选中通道的波形
            ax = figure.add_subplot(1, 1, 1)
            colors = {1: 'blue', 2: 'orange', 3: 'green', 4: 'red'}
            max_length = 0
            for channel in selected_channels:
                data = channel_data[channel]
                x = np.arange(len(data))  # x轴数据
                color = colors[channel]
                ax.plot(x, data, color=color, label=f'channel {channel}')
                max_length = max(max_length, len(data))

            ax.legend()
            y = [-2, -1, 0, 1, 2]
            ax.set_yticks(y)
            ax.set_yticklabels(y)
            ax.set_xlim(0, max_length - 1)  # 设置x轴从0到数据末端
            scene.addWidget(canvas)
            return X_kf
        except Exception as e:
            X_kf = []
            return X_kf

    def features(self, signal):
        x = signal
        X_rms2 = np.sum(x ** 2) / len(x)  # 3.均方值
        X_rms = np.sqrt(X_rms2)  # 4.均方根值(有效值)
        X_beta = np.mean(np.power(x, 4))  # 12.峭度
        X_kf = X_beta / X_rms ** 4  # 18.峭度指标
        X_kf = "{:.3f}".format(X_kf)
        return float(X_kf)

    def show_waveform(self):
        try:
            data = self.get_data()
            waveform_dialog = WaveformDialog(data)  # 替换为实际的文件路径
            selected_channels = [1, 2, 3, 4]  # 替换为根据需求选择的通道
            waveform_dialog.plot_waveform(selected_channels)
            waveform_dialog.exec_()
        except:
            pass

    def check_path(self):
        filename = self.file_content['filename']
        parent_path = self.file_content['parent_path']

        if 1 <= filename <= 9:
            csv_path = f"{parent_path}/0{filename}.csv"
        else:
            csv_path = f"{parent_path}/{filename}.csv"

        if os.path.exists(csv_path):
            self.load_text.setText(csv_path)
            self.csv_num_2.setText(str(filename))  # 显示csv序号

        return csv_path

    def check_photo(self):
        imgjpg_num = self.imgjpg_num.value()
        image_name = self.file_content["filename"] + imgjpg_num
        self.jpg_num.setText(str(image_name))  # 显示图片序号
        image_number = int(image_name)
        if image_number < 1:
            return None, None, image_name

        start = 0
        while True:
            end = start + 10000
            if start < image_number <= end:
                kind = f"{start + 1}-{end}"
                pic_path20 = f"{self.pic_path}/{kind}/20/{image_name}.jpg"
                pic_path21 = f"{self.pic_path}/{kind}/21/{image_name}.jpg"
                return pic_path20, pic_path21, image_name
            start += 10000

    def main(self):
        self.running()

    def csv_alert(self):
        msg_box = QMessageBox(self)
        msg_box.setWindowTitle("警告")
        msg_box.setText("峭度指标异常，请复核")
        msg_box.setIcon(QMessageBox.Warning)
        msg_box.setStandardButtons(QMessageBox.Cancel)

        # 创建一个定时器
        timer = QTimer(self)
        timer.setSingleShot(True)
        timer.timeout.connect(msg_box.close)
        timer.start(3000)  # 3秒后关闭

        msg_box.exec_()  # 使用 exec_() 而不是 show()

    def IoUslider_valuechange(self):  # 修改IoU的滑动条
        self.IoU_num.setValue(self.IoU_slider.value() / 100)

    def IoUnum_valuechange(self):  # 修改IoU的数字条
        self.IoU_slider.setValue(self.IoU_num.value() * 100)

    def confslider_valuechange(self):  # 修改conf的滑动条
        self.conf_num.setValue(self.conf_slider.value() / 100)

    def confnum_valuechange(self):  # 修改conf的数字条
        self.conf_slider.setValue(self.conf_num.value() * 100)

    def imgjpg_slider_valuechange(self):  # 修改csv的滑动条
        self.imgjpg_num.setValue(self.imgjpg_slider.value())

    def imgjpg_num_valuechange(self):  # 修改csv的数字条
        self.imgjpg_slider.setValue(self.imgjpg_num.value())

    def auto(self, checked):
        if checked:
            self.timer.start()  # 启动定时器
        else:
            self.timer.stop()

    def closeEvent(self, event):
        # 退出表示弹出框标题，"你确定退出吗？"表示弹出框的内容
        reply = QMessageBox.question(self, '确认退出', '您确定要退出吗？',
                                     QMessageBox.Yes | QMessageBox.No, QMessageBox.No)
        if reply == QMessageBox.Yes:
            event.accept()  # 接受关闭事件
        else:
            event.ignore()  # 忽略关闭事件


class WaveformDialog(QtWidgets.QDialog):
    def __init__(self, data):
        super().__init__()
        self.setWindowTitle("频谱图")
        self.setGeometry(100, 100, 800, 600)

        channel_1 = data.iloc[:, 2].values.astype(float)
        channel_2 = data.iloc[:, 5].values.astype(float)
        channel_3 = data.iloc[:, 8].values.astype(float)
        channel_4 = data.iloc[:, 11].values.astype(float)
        channel_1 = channel_1 * 1000 / (2.463 * 9.8)
        channel_2 = channel_2 * 1000 / (2.46 * 9.8)
        channel_3 = channel_3 * 1000 / (2.466 * 9.8)
        channel_4 = channel_4 * 1000 / (2.427 * 9.8)
        self.channel_data = {
            1: channel_1,  # 替换为通道1的实际数据列表
            2: channel_2,  # 替换为通道2的实际数据列表
            3: channel_3,  # 替换为通道3的实际数据列表
            4: channel_4  # 替换为通道4的实际数据列表
        }

        layout = QtWidgets.QVBoxLayout(self)
        self.figure = Figure(figsize=(18, 3.4), dpi=100)
        self.canvas = FigureCanvas(self.figure)
        layout.addWidget(self.canvas)

    def plot_waveform(self, selected_channels):
        self.figure.clear()
        ax = self.figure.add_subplot(1, 1, 1)
        colors = {1: 'blue', 2: 'orange', 3: 'green', 4: 'red'}
        legend_labels = []
        for channel in selected_channels:
            channel_data = self.channel_data[channel]
            fft_data = fft(channel_data)

            # 计算频率轴
            n = len(fft_data)
            timestep = 1  # 假设采样频率为1
            freq = np.fft.fftfreq(n, d=timestep)

            # 获取单边频谱
            half_n = n // 2
            freq = freq[:half_n]
            fft_data = np.abs(fft_data[:half_n])

            color = colors[channel]
            line, = ax.plot(freq, fft_data, color=color)
            legend_labels.append(line)

        ax.set_xlabel('Frequency')
        ax.set_ylabel('Amplitude')
        ax.set_xlim(0, 0.5)  # 设置x轴范围为0到0.5，表示正频率范围
        ax.set_ylim(0)  # 设置y轴范围为非负值

        ax.legend(legend_labels, selected_channels)  # 添加图例

        self.canvas.draw()


if __name__ == '__main__':
    from multiprocessing import freeze_support

    freeze_support()
    # 适应高DPI设备
    QtCore.QCoreApplication.setAttribute(QtCore.Qt.AA_EnableHighDpiScaling)
    # 适应Windows缩放
    QtGui.QGuiApplication.setAttribute(QtCore.Qt.HighDpiScaleFactorRoundingPolicy.PassThrough)

    app = QApplication(sys.argv)
    myWin = MyWindow()
    myWin.show()
    sys.exit(app.exec_())
